#!/bin/bash
# Linux脚本：批量导出和编辑Conda环境

# 创建输出目录
output_dir="conda_environments"
mkdir -p "$output_dir"

# 获取所有环境列表
envs=$(conda env list | grep -v "^#" | awk '{print $1}' | grep -v "^$")

echo "发现以下Conda环境："
for env in $envs; do
    echo "- $env"
done
echo ""

# 导出每个环境
for env in $envs; do
    echo "正在导出环境: $env"
    
    # 导出完整环境
    full_output_path="$output_dir/${env}_full.yml"
    conda env export --name "$env" > "$full_output_path"
    
    # 导出跨平台兼容版本
    compat_output_path="$output_dir/${env}_compatible.yml"
    conda env export --name "$env" --no-builds > "$compat_output_path"
    
    # 导出requirements.txt（包括conda和pip包）
    requirements_path="$output_dir/${env}_requirements.txt"
    echo "正在导出包到requirements.txt..."
    
    # 创建临时文件
    temp_file=$(mktemp)
    
    # 添加额外的索引URL到临时文件顶部
    echo "--extra-index-url https://mirrors.aliyun.com/pypi/simple/" > "$temp_file"
    echo "--extra-index-url https://pypi.tuna.tsinghua.edu.cn/simple/" >> "$temp_file"
    
    # 获取conda包列表（包括python版本）
    conda list --name "$env" | grep -v "^#" | while read -r line; do
        if [ -n "$line" ]; then
            package_name=$(echo "$line" | awk '{print $1}')
            version=$(echo "$line" | awk '{print $2}')
            
            # 排除一些不需要的系统包
            exclude_packages=("_libgcc_mutex" "_openmp_mutex" "libgcc" "libgomp" "libstdcxx" "ld_impl_linux-64" "libffi")
            skip=0
            for exclude in "${exclude_packages[@]}"; do
                if [ "$package_name" == "$exclude" ]; then
                    skip=1
                    break
                fi
            done
            
            if [ $skip -eq 0 ]; then
                echo "$package_name==$version" >> "$temp_file"
            fi
        fi
    done
    
    # 使用pip freeze获取pip包
    source activate "$env" && pip freeze >> "$temp_file"
    
    if [ -f "$temp_file" ]; then
        # 处理requirements.txt文件，确保格式为package==version或package>=version
        processed_temp_file=$(mktemp)
        
        # 添加额外的索引URL到处理后的临时文件顶部
        echo "--extra-index-url https://mirrors.aliyun.com/pypi/simple/" > "$processed_temp_file"
        echo "--extra-index-url https://pypi.tuna.tsinghua.edu.cn/simple/" >> "$processed_temp_file"
        
        # 处理每一行并去除重复的包
        declare -A unique_packages
        
        while IFS= read -r line; do
            # 跳过空行、注释和已添加的索引URL
            if [[ -z "$line" || "$line" =~ ^# || "$line" =~ ^--extra-index-url ]]; then
                continue
            fi
            
            # 处理已经符合格式的行
            if [[ "$line" =~ ^([a-zA-Z0-9_\-\.]+)(==|>=)(.+)$ ]]; then
                package_name="${BASH_REMATCH[1],,}" # 转换为小写
                if [ -z "${unique_packages[$package_name]}" ]; then
                    unique_packages[$package_name]=1
                    echo "$line" >> "$processed_temp_file"
                fi
                continue
            fi
            
            # 处理需要转换的行
            if [[ "$line" =~ ^([a-zA-Z0-9_\-\.]+)=(.+)$ ]]; then
                package_name="${BASH_REMATCH[1],,}" # 转换为小写
                version="${BASH_REMATCH[2]}"
                if [ -z "${unique_packages[$package_name]}" ]; then
                    unique_packages[$package_name]=1
                    echo "$package_name==$version" >> "$processed_temp_file"
                fi
            elif [[ "$line" =~ ^([a-zA-Z0-9_\-\.]+)$ ]]; then
                package_name="${BASH_REMATCH[1],,}" # 转换为小写
                if [ -z "${unique_packages[$package_name]}" ]; then
                    unique_packages[$package_name]=1
                    # 对于没有版本的包，使用>=0.0.0格式
                    echo "$package_name>=0.0.0" >> "$processed_temp_file"
                fi
            else
                # 保留其他格式的行
                echo "$line" >> "$processed_temp_file"
            fi
        done < "$temp_file"
        
        # 将处理后的内容移回原文件
        mv "$processed_temp_file" "$requirements_path"
        
        echo "Requirements.txt文件已格式化为package==version或package>=version格式。"
        
        # 检查是否有内容
        content_count=$(grep -v "^--" "$requirements_path" | grep -v "^$" | wc -l)
        if [ "$content_count" -eq 0 ]; then
            echo "警告: 环境 $env 中未找到包。至少添加Python版本。" >&2
            
            # 获取Python版本
            python_version=$(conda list --name "$env" python | grep -v "^#" | awk '$1=="python" {print $2}')
            if [ -n "$python_version" ]; then
                echo "python==$python_version" >> "$requirements_path"
            fi
        fi
    else
        echo "警告: 无法为环境 $env 创建requirements.txt文件" >&2
    fi
    
    # 清理临时文件
    if [ -f "$temp_file" ]; then
        rm "$temp_file"
    fi
    if [ -f "$processed_temp_file" ]; then
        rm "$processed_temp_file"
    fi
    
    echo "已导出到: $full_output_path, $compat_output_path 和 $requirements_path"
    echo "--------------------------------------------------"
done

echo -e "\n所有环境已导出到 $output_dir 目录"
echo "您可以使用以下命令重建环境："
echo "1. 从YAML文件: conda env create -f [环境文件路径].yml"
echo "2. 从requirements.txt:"
echo "   conda create -n [环境名称] python=[版本]"
echo "   conda activate [环境名称]"
echo "   pip install -r [requirements文件路径].txt" 